package com.zyb.lmbackend.llm;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.base.Stopwatch;
import com.zyb.lmbackend.config.AiProperties;
import com.zyb.lmbackend.constants.LlmPrompts;
import com.zyb.lmbackend.entity.LlmCallAudit;
import com.zyb.lmbackend.llm.model.CategoryDef;
import com.zyb.lmbackend.llm.model.ClsResult;
import com.zyb.lmbackend.service.LlmCallAuditService;
import com.zyb.lmbackend.util.InvalidLlmOutputException;
import com.zyb.lmbackend.util.JsonLinesOrArrayParser;
import com.zyb.lmbackend.util.JsonSchemaValidatorUtil;
import io.micrometer.core.instrument.MeterRegistry;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.client.RestTemplate;
import org.springframework.http.client.SimpleClientHttpRequestFactory;

/**
 * 直连 DashScope 文本生成接口的 LLM 客户端实现。
 * <p>
 * 通过 HTTP 调用 generation 接口，手工构造 messages 与参数，解析返回 JSON。
 * 支持可配置超时，适合无需 Agent/App 模式的直接调用场景。
 */
@RequiredArgsConstructor
@Slf4j
public class DirectDashscopeLlmClient implements LlmClient {

    private final AiProperties aiProperties;
    private final LlmCallAuditService auditService;
    private final MeterRegistry meterRegistry;
    private final ObjectMapper objectMapper = new ObjectMapper();

    private static final String CHAT_COMPLETIONS_ENDPOINT = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation";

    /** 创建带连接/读超时的 RestTemplate 实例。 */
    private RestTemplate newHttp(int connectTimeoutMs, int readTimeoutMs) {
        SimpleClientHttpRequestFactory f = new SimpleClientHttpRequestFactory();
        f.setConnectTimeout(connectTimeoutMs);
        f.setReadTimeout(readTimeoutMs);
        return new RestTemplate(f);
    }

    /** 将如 900s/10m/600000 等超时配置解析为毫秒数。 */
    private int parseTimeoutMs(String val, int defaultMs) {
        if (val == null || val.isBlank()) return defaultMs;
        String s = val.trim().toLowerCase();
        try {
            if (s.endsWith("ms")) {
                return Math.max(0, Integer.parseInt(s.substring(0, s.length() - 2)));
            } else if (s.endsWith("s")) {
                return Math.max(0, (int) (Double.parseDouble(s.substring(0, s.length() - 1)) * 1000));
            } else if (s.endsWith("m")) {
                return Math.max(0, (int) (Double.parseDouble(s.substring(0, s.length() - 1)) * 60_000));
            } else if (s.startsWith("pt") || s.startsWith("pT") || s.startsWith("Pt")) {
                // basic ISO-8601 duration support like PT90S
                java.time.Duration d = java.time.Duration.parse(val);
                long ms = d.toMillis();
                return (ms > Integer.MAX_VALUE) ? Integer.MAX_VALUE : (int) ms;
            } else {
                // plain number -> milliseconds
                long ms = Long.parseLong(s);
                return (ms > Integer.MAX_VALUE) ? Integer.MAX_VALUE : (int) ms;
            }
        } catch (Exception ignore) {
            return defaultMs;
        }
    }

    @Override
    /** 调用大模型对样本集合进行类目归纳。 */
    public List<CategoryDef> buildTaxonomy(List<String> samples) {
        var cfg = aiProperties.taxonomyCfg();
        if (cfg == null) throw new IllegalStateException("Direct model 未配置 taxonomy");
        String systemPrompt = LlmPrompts.TAXONOMY_SYSTEM_PROMPT;
        String content = toJsonSafe(new HashMap<String, Object>() {{ put("samples", samples); }});

        Map<String, Object> req = new HashMap<>();
        req.put("model", cfg.getModel());
        req.put("input", Map.of("messages", List.of(
                Map.of("role", "system", "content", systemPrompt),
                Map.of("role", "user", "content", content)
        )));
        Map<String, Object> params = new HashMap<>();
        if (cfg.getTemperature() != null) params.put("temperature", cfg.getTemperature());
        params.put("result_format", "message");
        params.put("incremental_output", false);
        req.put("parameters", params);

        String inJson = toJsonSafe(req);

        Stopwatch sw = Stopwatch.createStarted();
        Integer status = null;
        Integer tokens = null;
        String outText = null;
        try {
            int connectMs = 10_000;
            int readMs = parseTimeoutMs(cfg.getTimeout(), 900_000); // taxonomy can be slow (default 15min)
            //这里调用给大模型
            String raw = callDirect(cfg.getApiKey(), inJson, connectMs, readMs);
            status = 200; // if no exception
            tokens = extractTokens(raw);
            outText = extractContent(raw);

            List<CategoryDef> result = objectMapper.readValue(outText, new TypeReference<List<CategoryDef>>() {});
            return result;
        } catch (Exception e) {
            meterRegistry.counter("llm.errors").increment();
            log.error("Parse taxonomy output failed: {}", truncate64k(outText));
            throw new InvalidLlmOutputException("无法解析类目归纳结果", e);
        } finally {
            long latency = sw.stop().elapsed(TimeUnit.MILLISECONDS);
            meterRegistry.counter("llm.calls").increment();
            meterRegistry.timer("llm.latency").record(latency, TimeUnit.MILLISECONDS);
            saveAudit(null, null, "IN", truncate64k(inJson), tokens, null, null);
            saveAudit(null, null, "OUT", truncate64k(outText), tokens, (int) latency, status);
        }
    }

    @Override
    /**
     * 对一批文本进行分类。
     * @param details 需要分类的文本列表
     * @param allowedCodes 允许输出的类目 code 列表
     */
    public List<ClsResult> classifyBatch(List<String> details, List<String> allowedCodes) {
        var cfg = aiProperties.classifyCfg();
        if (cfg == null) throw new IllegalStateException("Direct model 未配置 classify");
        String sys = LlmPrompts.CLASSIFY_SYSTEM_PROMPT + "\n允许的code列表：" + String.join(",", allowedCodes);
        String user = toJsonSafe(details);

        Map<String, Object> req = new HashMap<>();
        req.put("model", cfg.getModel());
        req.put("input", Map.of("messages", List.of(
                Map.of("role", "system", "content", sys),
                Map.of("role", "user", "content", user)
        )));
        Map<String, Object> params = new HashMap<>();
        if (cfg.getTemperature() != null) params.put("temperature", cfg.getTemperature());
        params.put("result_format", "message");
        params.put("incremental_output", false);
        req.put("parameters", params);

        String inJson = toJsonSafe(req);

        Stopwatch sw = Stopwatch.createStarted();
        Integer status = null;
        Integer tokens = null;
        String raw = null;
        try {
            int connectMs = 10_000;
            int readMs = parseTimeoutMs(cfg.getTimeout(), 900_000); // default 15min for safety
            raw = callDirect(cfg.getApiKey(), inJson, connectMs, readMs);
            status = 200;
            tokens = extractTokens(raw);
            String content = extractContent(raw);

            JsonLinesOrArrayParser parser = new JsonLinesOrArrayParser(objectMapper);
            List<JsonNode> nodes = parser.parseToNodes(content);
            new JsonSchemaValidatorUtil(objectMapper).validateEach(nodes, "schema/classify_item.schema.json");
            List<ClsResult> list = new ArrayList<>();
            for (JsonNode n : nodes) {
                int index = n.path("index").asInt();
                String code = n.path("category_code").asText();
                double conf = n.path("confidence").asDouble();
                String reason = n.path("reason").asText();
                list.add(new ClsResult(index, code, conf, reason));
            }
            return list;
        } catch (Exception e) {
            meterRegistry.counter("llm.errors").increment();
//            log.error("Parse classify output failed: {}", truncate64k(raw));
            throw new InvalidLlmOutputException("无法解析分类结果", e);
        } finally {
            long latency = sw.stop().elapsed(TimeUnit.MILLISECONDS);
            meterRegistry.counter("llm.calls").increment();
            meterRegistry.timer("llm.latency").record(latency, TimeUnit.MILLISECONDS);
            saveAudit(null, null, "IN", truncate64k(inJson), tokens, null, null);
            saveAudit(null, null, "OUT", truncate64k(raw), tokens, (int) latency, status);
        }
    }

    /** 以 Bearer Token 直连 DashScope generation 接口。 */
    private String callDirect(String apiKey, String bodyJson, int connectTimeoutMs, int readTimeoutMs) {
        RestTemplate http = newHttp(connectTimeoutMs, readTimeoutMs);
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.setBearerAuth(apiKey);
        HttpEntity<String> entity = new HttpEntity<>(bodyJson, headers);
        ResponseEntity<String> resp = http.postForEntity(CHAT_COMPLETIONS_ENDPOINT, entity, String.class);
        return resp.getBody() == null ? "" : resp.getBody();
    }

    /** 从 DashScope 返回 JSON 提取首个 message.content 文本。 */
    private String extractContent(String raw) throws JsonProcessingException {
        if (raw == null) return "";
        JsonNode root = objectMapper.readTree(raw);
        // Try typical paths for dashscope text-generation
        JsonNode content = null;
        if (root.has("output")) {
            JsonNode out = root.get("output");
            if (out.has("choices")) {
                JsonNode arr = out.get("choices");
                if (arr.isArray() && arr.size() > 0) {
                    JsonNode ch0 = arr.get(0);
                    if (ch0.has("message")) {
                        content = ch0.get("message").get("content");
                    }
                }
            }
            if (content == null && out.has("text")) content = out.get("text");
        }
        if (content == null && root.has("choices")) {
            JsonNode arr = root.get("choices");
            if (arr.isArray() && arr.size() > 0) {
                JsonNode ch0 = arr.get(0);
                if (ch0.has("message")) content = ch0.get("message").get("content");
            }
        }
        return content == null || content.isNull() ? "" : content.asText();
    }

    private Integer extractTokens(String raw) {
        try {
            JsonNode root = objectMapper.readTree(raw);
            JsonNode usage = root.path("usage");
            if (usage.isMissingNode() || usage.isNull()) return null;
            int in = 0, out = 0;
            if (usage.has("total_tokens")) {
                return usage.get("total_tokens").asInt();
            }
            if (usage.has("input_tokens")) in = usage.get("input_tokens").asInt();
            if (usage.has("output_tokens")) out = usage.get("output_tokens").asInt();
            // sometimes nested in models[0]
            if (in == 0 && out == 0 && usage.has("models")) {
                JsonNode arr = usage.get("models");
                if (arr.isArray() && arr.size() > 0) {
                    JsonNode m0 = arr.get(0);
                    if (m0.has("input_tokens")) in = m0.get("input_tokens").asInt();
                    if (m0.has("output_tokens")) out = m0.get("output_tokens").asInt();
                }
            }
            int sum = in + out;
            return sum > 0 ? sum : null;
        } catch (Exception ignore) {
            return null;
        }
    }

    private void saveAudit(Long taskId, Integer batchNo, String direction, String payload,
            Integer tokens, Integer latencyMs, Integer httpStatus) {
        LlmCallAudit audit = new LlmCallAudit();
        audit.setTaskId(taskId);
        audit.setBatchNo(batchNo);
        audit.setDirection(direction);
        audit.setPayload(truncate64k(payload));
        audit.setTokens(tokens);
        audit.setLatencyMs(latencyMs);
        audit.setHttpStatus(httpStatus);
        auditService.save(audit);
    }

    private String toJsonSafe(Object obj) {
        try {
            return objectMapper.writeValueAsString(obj);
        } catch (JsonProcessingException e) {
            return String.valueOf(obj);
        }
    }

    private String truncate64k(String s) {
        if (s == null) return null;
        byte[] bytes = s.getBytes(StandardCharsets.UTF_8);
        int limit = 64 * 1024; // 64KB
        if (bytes.length <= limit) return s;
        int end = limit;
        while (end > 0 && (bytes[end - 1] & 0xC0) == 0x80) {
            end--;
        }
        return new String(bytes, 0, end, StandardCharsets.UTF_8);
    }
}
